#!/usr/bin/env python
# -*- coding: utf-8 -*-

from flask import Flask, render_template, request, jsonify, redirect, url_for, session, flash
import networkx as nx
import json
import jieba
import re
from pyvis.network import Network
import os
import pandas as pd
import numpy as np
import uuid
from werkzeug.security import generate_password_hash, check_password_hash
from werkzeug.utils import secure_filename
from functools import wraps
import io
import PyPDF2
try:
    import docx
except ImportError:
    docx = None

app = Flask(__name__)
app.secret_key = os.urandom(24)  # 用于session加密

# 文件上传设置
UPLOAD_FOLDER = os.path.join('static', 'uploads')
ALLOWED_EXTENSIONS = {'txt', 'pdf', 'doc', 'docx'}
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 最大上传16MB

# 确保上传目录存在
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# 模拟用户数据库
users = {
    # 'username': {'password_hash': '...', 'name': '...'}
}

# 全局变量存储知识图谱
knowledge_graph = nx.DiGraph()
# 存储已知实体和关系 - 预置实体和关系列表
known_entities = ['牛顿', '爱因斯坦', '居里夫人', '特斯拉', 
               '相对论', '万有引力', '量子力学', 'E=mc²', 
               '麦克斯韦方程组', '光子', '质能方程', '电磁学',
               '放射性', '镭', '波兰', '法国', '诺贝尔奖',
               '电磁感应', '交流电', '无线电', '电动机',
               '剑桥大学', '物理学', '数学', '力学',
               '普林斯顿大学', '索尔维会议', '曼哈顿计划']
known_relations = ['发明', '发现', '提出', '证明', '研究', '获得', '就读于', '工作于', 
                '出生于', '组成', '包含', '应用于', '推导出', '参与', '合作', '指导']
# 存储新发现的三元组
new_triples = []

# 预置的演示文本
demo_texts = [
    "牛顿是万有引力定律的提出者，他在剑桥大学研究物理学和数学，并发明了微积分。",
    "爱因斯坦提出了相对论，并推导出了著名的质能方程E=mc²。他曾在普林斯顿大学工作，并参与了索尔维会议。",
    "居里夫人出生于波兰，但在法国工作。她发现了放射性元素镭，并因此获得了诺贝尔奖。",
    "特斯拉在电磁学领域有重要贡献，他发明了交流电系统和无线电技术，电动机的设计也与他有关。",
    "量子力学是现代物理学的重要分支，它与相对论共同组成了现代物理学的两大基本理论。爱因斯坦虽然提出了光子概念，但他对量子力学的某些解释持怀疑态度。",
    "麦克斯韦方程组是电磁学的基础，它证明了电和磁是同一种现象的不同表现形式。这一理论对特斯拉的发明有重要影响。"
]

# 登录验证装饰器
def login_required(f):
    @wraps(f)
    def decorated_function(*args, **kwargs):
        if 'username' not in session:
            flash('请先登录', 'danger')
            return redirect(url_for('login'))
        return f(*args, **kwargs)
    return decorated_function

def extract_triples(text, entities, relations):
    """
    从文本中提取三元组
    
    Args:
        text: 输入文本
        entities: 已知实体列表
        relations: 已知关系列表
        
    Returns:
        发现的三元组列表 [(head, relation, tail), ...]
    """
    # 使用jieba进行分词
    words = jieba.lcut(text)
    
    # 找到文本中的实体
    found_entities = []
    for entity in entities:
        if entity in text:
            found_entities.append(entity)
    
    triples = []
    # 简单的三元组提取逻辑：遍历所有实体对，检查它们之间是否存在关系
    for head in found_entities:
        for tail in found_entities:
            if head != tail:
                head_pos = text.find(head)
                tail_pos = text.find(tail)
                
                # 获取两个实体之间的文本片段
                if head_pos < tail_pos:
                    middle_text = text[head_pos + len(head):tail_pos]
                else:
                    middle_text = text[tail_pos + len(tail):head_pos]
                
                # 检查关系是否出现在中间文本中
                for relation in relations:
                    if relation in middle_text:
                        if head_pos < tail_pos:
                            triples.append((head, relation, tail))
                        else:
                            triples.append((tail, relation, head))
    
    return triples

def visualize_graph(graph, new_triples=None):
    """生成知识图谱的可视化HTML文件"""
    # 设置高级配置以美化图表
    net = Network(height='750px', width='100%', bgcolor='#f9f9f9', font_color='#333')
    net.set_options("""
    {
      "nodes": {
        "shape": "dot",
        "size": 25,
        "font": {
          "size": 14,
          "face": "Microsoft YaHei", 
          "bold": false,
          "color": "#333"
        },
        "borderWidth": 2,
        "shadow": true,
        "scaling": {
          "label": {
            "min": 12,
            "max": 22
          }
        }
      },
      "edges": {
        "width": 1.5,
        "smooth": {
          "type": "continuous",
          "forceDirection": "none",
          "roundness": 0.55
        },
        "font": {
          "size": 13, 
          "face": "Microsoft YaHei",
          "background": "rgba(255, 255, 255, 0.7)",
          "color": "#333"
        },
        "shadow": true,
        "arrows": {
          "to": {
            "enabled": true,
            "scaleFactor": 0.8
          }
        }
      },
      "physics": {
        "barnesHut": {
          "gravitationalConstant": -5000,
          "centralGravity": 0.3,
          "springLength": 180,
          "springConstant": 0.05,
          "damping": 0.09
        },
        "minVelocity": 0.75,
        "solver": "barnesHut"
      },
      "interaction": {
        "navigationButtons": true,
        "keyboard": true,
        "hover": true,
        "multiselect": true,
        "tooltipDelay": 200
      }
    }
    """)
    
    # 节点类型分组，根据节点特征进行分类
    entity_types = {
        '人物': ['牛顿', '爱因斯坦', '居里夫人', '特斯拉', '物理学家', '数学家'],
        '概念': ['相对论', '万有引力', '量子力学', 'E=mc²', '麦克斯韦方程组', '光子', '质能方程', '电磁学', '放射性', '物理学', '数学', '力学', '交流电', '无线电'],
        '机构': ['剑桥大学', '普林斯顿大学', '索尔维会议', '曼哈顿计划'],
        '地点': ['波兰', '法国'],
        '奖项': ['诺贝尔奖'],
        '元素': ['镭'],
        '设备': ['电动机']
    }
    
    # 为不同类型节点分配不同颜色
    colors = {
        '人物': '#4e79a7',
        '概念': '#f28e2c',
        '机构': '#e15759',
        '地点': '#76b7b2',
        '奖项': '#59a14f',
        '元素': '#af7aa1',
        '设备': '#ff9da7'
    }
    
    # 添加节点
    for node in graph.nodes():
        # 确定节点类型
        node_type = None
        for t, entities in entity_types.items():
            if node in entities:
                node_type = t
                break
        
        # 如果没有找到对应类型，默认为'其他'
        if node_type is None:
            node_type = '其他'
            color = '#b07aa1'
        else:
            color = colors[node_type]
        
        # 创建丰富的工具提示内容
        connections = list(graph.neighbors(node))
        tooltip = f"""<div style="padding:10px;max-width:300px;font-family:'Microsoft YaHei',sans-serif;">
            <h4 style="margin:0 0 5px 0;color:#333;">{node}</h4>
            <p style="margin:5px 0;"><strong>类型:</strong> {node_type}</p>
            <p style="margin:5px 0;"><strong>连接数:</strong> {len(connections)}</p>"""
        
        # 如果有关联节点，添加到工具提示中
        if connections:
            tooltip += "<p style='margin:5px 0;'><strong>关联节点:</strong></p><ul style='margin:5px 0;padding-left:20px;'>"
            for neighbor in connections[:5]:  # 只显示前5个关联
                edge_data = graph.get_edge_data(node, neighbor) or graph.get_edge_data(neighbor, node) or {}
                relation = edge_data.get('relation', '关联')
                tooltip += f"<li>{neighbor} ({relation})</li>"
            if len(connections) > 5:
                tooltip += f"<li>...等共 {len(connections)} 个关联</li>"
            tooltip += "</ul>"
        
        tooltip += "</div>"
        
        # 添加节点并设置样式
        net.add_node(
            node, 
            label=node, 
            title=tooltip, 
            color=color,
            size=35 if node == '牛顿' else (25 if len(connections) > 5 else 20),  # 根据重要性和连接度调整大小
            borderWidth=2,
            shape='star' if node == '牛顿' else ('box' if node_type == '著作' else ('diamond' if node_type == '实验' else 'dot')),
            shadow=True,
            mass=2 if node == '牛顿' else 1  # 影响物理引擎中的稳定性
        )
    
    # 添加边，新的三元组用特殊颜色标记
    for u, v, data in graph.edges(data=True):
        relation = data.get('relation', '')
        is_new = new_triples is not None and (u, relation, v) in new_triples
        
        # 设置边的样式
        color = '#e41a1c' if is_new else '#999'
        width = 2.5 if is_new else 1.5
        dashes = True if relation in ['发现', '提出', '发明'] else False
        
        # 创建边的丰富工具提示
        edge_tooltip = f"""<div style="padding:8px;max-width:250px;font-family:'Microsoft YaHei',sans-serif;">
            <p style="margin:2px 0;font-size:14px;"><strong>{u}</strong> {relation} <strong>{v}</strong></p>
            <p style="margin:2px 0;color:{'#e41a1c' if is_new else '#666'};font-size:12px;">
                {('★ 新发现关系' if is_new else '已知关系')}
            </p>
        </div>"""
        
        # 添加关系边
        net.add_edge(
            u, v,
            title=edge_tooltip,
            label=relation,
            color=color,
            width=width,
            dashes=dashes,
            physics=True,
            smooth={'type': 'curvedCW', 'roundness': 0.2} if is_new else {'type': 'continuous'},
            font={
                'size': 12,
                'color': '#333',
                'background': 'rgba(255,255,255,0.7)',
                'strokeWidth': 0
            },
            shadow=is_new,  # 为新关系添加阴影效果
            value=2 if is_new else 1  # 影响边的视觉权重
        )
    
    # 保存为HTML文件
    net.save_graph(os.path.join('static', 'js', 'graph.html'))
    return

@app.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        username = request.form.get('username')
        password = request.form.get('password')
        
        if username in users and check_password_hash(users[username]['password_hash'], password):
            session['username'] = username
            session['name'] = users[username].get('name', username)
            flash('登录成功', 'success')
            return redirect(url_for('index'))
        else:
            flash('用户名或密码错误', 'danger')
    
    return render_template('login.html')

@app.route('/register', methods=['GET', 'POST'])
def register():
    if request.method == 'POST':
        username = request.form.get('username')
        password = request.form.get('password')
        name = request.form.get('name')
        
        if username in users:
            flash('用户名已存在', 'danger')
            return render_template('register.html')
        
        users[username] = {
            'password_hash': generate_password_hash(password),
            'name': name
        }
        
        flash('注册成功，请登录', 'success')
        return redirect(url_for('login'))
    
    return render_template('register.html')

@app.route('/logout')
def logout():
    session.pop('username', None)
    session.pop('name', None)
    flash('已退出登录', 'info')
    return redirect(url_for('login'))

@app.route('/', methods=['GET'])
@login_required
def index():
    return render_template('index.html', 
                          known_entities=known_entities, 
                          known_relations=known_relations, 
                          new_triples=new_triples,
                          name=session.get('name', ''))

@app.route('/setup', methods=['POST'])
@login_required
def setup():
    global known_entities, known_relations, knowledge_graph, new_triples
    
    # 获取表单数据
    entities_text = request.form.get('entities', '')
    relations_text = request.form.get('relations', '')
    
    # 处理实体和关系列表
    known_entities = [e.strip() for e in entities_text.split(',') if e.strip()]
    known_relations = [r.strip() for r in relations_text.split(',') if r.strip()]
    
    # 重置图形和新三元组
    knowledge_graph = nx.DiGraph()
    new_triples = []
    
    # 添加已知实体作为节点
    for entity in known_entities:
        knowledge_graph.add_node(entity)
    
    # 生成可视化
    visualize_graph(knowledge_graph)
    
    return redirect(url_for('index'))

@app.route('/extract', methods=['POST'])
@login_required
def extract():
    global knowledge_graph, new_triples
    
    # 获取文本
    text = request.form.get('text', '')
    
    # 提取三元组
    extracted_triples = extract_triples(text, known_entities, known_relations)
    
    # 添加到图中，并记录新添加的三元组
    added_triples = []
    for head, relation, tail in extracted_triples:
        # 检查是否是新的三元组
        if not knowledge_graph.has_edge(head, tail) or knowledge_graph[head][tail].get('relation') != relation:
            knowledge_graph.add_edge(head, tail, relation=relation)
            added_triples.append((head, relation, tail))
    
    # 更新新发现的三元组列表
    new_triples = added_triples
    
    # 如果发现了新的三元组
    if added_triples:
        flash(f'成功提取{len(added_triples)}个三元组', 'success')
    else:
        flash('未发现新的三元组', 'info')
    
    # 更新可视化
    visualize_graph(knowledge_graph, new_triples)
    
    return redirect(url_for('index'))

@app.route('/clear', methods=['POST'])
@login_required
def clear():
    global knowledge_graph, new_triples
    
    # 清空图和新三元组
    knowledge_graph = nx.DiGraph()
    new_triples = []
    
    # 添加已知实体作为节点
    for entity in known_entities:
        knowledge_graph.add_node(entity)
    
    # 更新可视化
    visualize_graph(knowledge_graph)
    
    return redirect(url_for('index'))

@app.route('/load_demo', methods=['POST'])
@login_required
def load_demo():
    global knowledge_graph, new_triples
    
    # 清空图和新三元组
    knowledge_graph = nx.DiGraph()
    new_triples = []
    
    # 添加所有已知实体作为节点
    for entity in known_entities:
        knowledge_graph.add_node(entity)
    
    # 预置一些基本关系作为初始图谱
    initial_triples = [
        ('牛顿', '提出', '万有引力'),
        ('牛顿', '研究', '物理学'),
        ('牛顿', '研究', '数学'),
        ('牛顿', '就读于', '剑桥大学'),
        ('爱因斯坦', '提出', '相对论'),
        ('爱因斯坦', '推导出', 'E=mc²'),
        ('爱因斯坦', '工作于', '普林斯顿大学'),
        ('爱因斯坦', '参与', '索尔维会议'),
        ('居里夫人', '出生于', '波兰'),
        ('居里夫人', '工作于', '法国'),
        ('居里夫人', '发现', '镭'),
        ('居里夫人', '获得', '诺贝尔奖'),
        ('特斯拉', '发明', '交流电'),
        ('特斯拉', '发明', '无线电'),
        ('特斯拉', '研究', '电磁学'),
        ('量子力学', '组成', '物理学'),
        ('相对论', '组成', '物理学'),
    ]
    
    # 添加初始三元组到图中
    for head, relation, tail in initial_triples:
        if head in known_entities and tail in known_entities:
            knowledge_graph.add_edge(head, tail, relation=relation)
    
    # 更新可视化
    visualize_graph(knowledge_graph)
    
    # 设置演示文本，将其作为 flash 消息返回，供用户使用
    demo_text_message = "演示文本已加载，您可以复制使用这些文本进行实体关系提取：<br>"
    for i, text in enumerate(demo_texts, 1):
        demo_text_message += f"<strong>文本{i}:</strong> {text}<br>"
    
    flash(demo_text_message, 'info')
    return redirect(url_for('index'))

@app.route('/get_triples', methods=['GET'])
@login_required
def get_triples():
    global knowledge_graph
    
    # 收集所有三元组
    all_triples = []
    for u, v, data in knowledge_graph.edges(data=True):
        relation = data.get('relation', '')
        all_triples.append([u, relation, v])
    
    return jsonify(triples=all_triples, new_triples=new_triples)

def read_file_content(file):
    """从不同类型的文件中提取文本"""
    filename = secure_filename(file.filename)
    content = ""
    
    # 检查文件扩展名
    if filename.endswith('.txt'):
        content = file.read().decode('utf-8', errors='ignore')
    elif filename.endswith('.pdf'):
        try:
            pdf_reader = PyPDF2.PdfReader(io.BytesIO(file.read()))
            for page in pdf_reader.pages:
                content += page.extract_text() + "\n"
        except Exception as e:
            flash(f"PDF解析错误: {str(e)}", "danger")
            return ""
    elif filename.endswith(('.doc', '.docx')):
        if not docx:
            flash("系统未安装python-docx库，无法处理Word文档", "danger")
            return ""
        try:
            doc = docx.Document(io.BytesIO(file.read()))
            for para in doc.paragraphs:
                content += para.text + "\n"
        except Exception as e:
            flash(f"Word文档解析错误: {str(e)}", "danger")
            return ""
    else:
        flash("不支持的文件格式", "danger")
        return ""
    
    return content

@app.route('/extract_file', methods=['POST'])
@login_required
def extract_file():
    global knowledge_graph, new_triples
    
    if 'file' not in request.files:
        flash('未找到上传文件', 'danger')
        return redirect(url_for('index'))
    
    file = request.files['file']
    if file.filename == '':
        flash('未选择文件', 'danger')
        return redirect(url_for('index'))
    
    # 读取文件内容
    text = read_file_content(file)
    if not text:
        return redirect(url_for('index'))
    
    # 提取三元组
    extracted_triples = extract_triples(text, known_entities, known_relations)
    
    # 添加到图中，并记录新添加的三元组
    added_triples = []
    for head, relation, tail in extracted_triples:
        # 检查是否是新的三元组
        if not knowledge_graph.has_edge(head, tail) or knowledge_graph[head][tail].get('relation') != relation:
            knowledge_graph.add_edge(head, tail, relation=relation)
            added_triples.append((head, relation, tail))
    
    # 更新新发现的三元组列表
    new_triples = added_triples
    
    # 如果发现了新的三元组
    if added_triples:
        flash(f'从文件中成功提取{len(added_triples)}个三元组', 'success')
    else:
        flash('未从文件中发现新的三元组', 'info')
    
    # 更新可视化
    visualize_graph(knowledge_graph, new_triples)
    
    return redirect(url_for('index'))

if __name__ == '__main__':
    # 确保static/js目录存在
    os.makedirs(os.path.join('static', 'js'), exist_ok=True)
    
    # 创建示例用户（如果不存在）
    if 'admin' not in users:
        users['admin'] = {
            'password_hash': generate_password_hash('admin'),
            'name': '管理员'
        }
    
    app.run(debug=True, port=5001)
